from brian import *
from brian.library.random_processes import *
from brian.library.synapses import  *
import MoNeuron as Moto
import InNeuron as IN
from time import clock
from MyConstants import gSynE, gSynI
from matplotlib.backends.backend_pdf import PdfPages





class rybak(object):
    
    binsize = 30*ms
    stype = 'exp'
    def plot_hist(self,ratemonitor):
        bar(ratemonitor.times/ms,ratemonitor.rate,width=self.binsize/ms,color='w')
        
    def plot_spike(self,spikemon): 
        raster_plot(spikemon)
    
    def plot_simple(self,trace):
        plot(trace.times/ms,trace[0]/mV)
       
        
    def plot_all(self,n,fun,lst):
        fig=figure(n)
        for i in xrange(0,lst.__len__()):
            subplot(lst.__len__(),1,i+1)
            fun(lst[i])
            fig.axes[i].set_xlim(0,self.DUR/ms)
        
        return fig
    
    def reinit_monitors(self,lst):
        for i in xrange(0,lst.__len__()):
            lst[i].reinit()
        
    def __init__(self,Evar=1,rand=False):

        self.Evar=Evar
        self.rand=rand
        self.setup_neurons()
        self.setup_connections()
        self.setup_monitors()
        self.setup_network_operations()

        
    def setup_neurons(self):
        
    
    
        self.MNFactory=Moto.MnNeuronFactory()
        self.MN=self.MNFactory.init(42,self.stype,self.Evar,self.rand)
        self.INFactory=IN.InNeuronFactory()
        self.P=self.INFactory.init(40,40,120,40,self.stype,self.Evar,self.rand)
        print 'init'
        self.RG_E = self.INFactory.SG_RG.subgroup(20)
        self.RG_F = self.INFactory.SG_RG.subgroup(20)
        self.RG_F.drive = 0.45
        self.RG_E.drive = 0.51
        #self.RG_E.drive = TimedArray([0.5, 0, 0.5],start=0*ms,dt=5*second)
        
        self.PF_E = self.INFactory.SG_PF.subgroup(20)
        self.PF_F = self.INFactory.SG_PF.subgroup(20)
        #self.PF_F.drive = TimedArray([0.5, 0.95, 0.5],start=0*ms,dt=5*second)
        
        self.INRG_E = self.INFactory.SG_INRG.subgroup(20)
        self.INRG_F = self.INFactory.SG_INRG.subgroup(20)
        
        self.INPF_E = self.INFactory.SG_IN.subgroup(20)
        self.INPF_F = self.INFactory.SG_IN.subgroup(20)
        
        self.IA_E = self.INFactory.SG_IN.subgroup(20)
        self.IA_F = self.INFactory.SG_IN.subgroup(20)
        
        self.R_E = self.INFactory.SG_IN.subgroup(20)
        self.R_F = self.INFactory.SG_IN.subgroup(20)
        
        self.MN_E = self.MN.subgroup(21)
        self.MN_F = self.MN.subgroup(21)
        
        self.realMN_E = self.MN_E.subgroup(20)
        self.fakeMN_E = self.MN_E.subgroup(1)
        
        self.realMN_F = self.MN_F.subgroup(20)
        self.fakeMN_F = self.MN_F.subgroup(1)
        
        
        
        self.net = Network(self.P,self.MN)
        
    def setup_connections(self,use_delay=False):
        #Connections
        self.RG_E_INRG_F=Connection(self.RG_E,self.INRG_F,'ge',weight=gSynE*0.45,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.RG_F_INRG_E=Connection(self.RG_F,self.INRG_E,'ge',weight=gSynE*0.45,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.RG_E_INRG_F,self.RG_F_INRG_E)
        
        self.RG_E_self=Connection(self.RG_E,self.RG_E,'ge',weight=gSynE*0.0125,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:2*ms+randn(20)*0.3*ms if use_delay else 0*ms);
        self.RG_F_self=Connection(self.RG_F,self.RG_F,'ge',weight=gSynE*0.0125,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:2*ms+randn(20)*0.3*ms if use_delay else 0*ms);
        self.net.add(self.RG_E_self,self.RG_F_self)
        
        self.RG_E_RG_F=Connection(self.RG_E,self.RG_F,'ge',weight=gSynE*0.0125,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.RG_F_RG_E=Connection(self.RG_F,self.RG_E,'ge',weight=gSynE*0.0125,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.RG_E_RG_F,self.RG_F_RG_E)
        
        self.INRG_E_RG_E=Connection(self.INRG_E,self.RG_E,'gi',weight=gSynI*0.115,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.INRG_F_RG_E=Connection(self.INRG_F,self.RG_F,'gi',weight=gSynI*0.115,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.INRG_E_RG_E,self.INRG_F_RG_E)
        
        self.RG_E_PF_E=Connection(self.RG_E,self.PF_E,'ge',weight=gSynE*0.0075,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.RG_F_PF_F=Connection(self.RG_F,self.PF_F,'ge',weight=gSynE*0.0075,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.RG_E_PF_E,self.RG_F_PF_F)
        
        self.INRG_E_PF_E=Connection(self.INRG_E,self.PF_E,'gi',weight=gSynI*0.05,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.INRG_F_PF_E=Connection(self.INRG_F,self.PF_F,'gi',weight=gSynI*0.05,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.INRG_E_PF_E,self.INRG_F_PF_E)
        
        self.INPF_E_PF_E=Connection(self.INPF_E,self.PF_E,'gi',weight=gSynI*0.35,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.INPF_F_PF_E=Connection(self.INPF_F,self.PF_F,'gi',weight=gSynI*0.35,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.INPF_E_PF_E,self.INPF_F_PF_E)
        
        self.PF_E_INPF_F=Connection(self.PF_E,self.INPF_F,'ge',weight=gSynE*0.2,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.PF_F_INPF_E=Connection(self.PF_F,self.INPF_E,'ge',weight=gSynE*0.2,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.PF_E_INPF_F,self.PF_F_INPF_E)
        
        self.PF_E_IA_E=Connection(self.PF_E,self.IA_E,'ge',weight=gSynE*0.4,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.PF_F_IA_F=Connection(self.PF_F,self.IA_F,'ge',weight=gSynE*0.4,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.PF_E_IA_E,self.PF_F_IA_F)
        
        self.IA_F_IA_E=Connection(self.IA_F,self.IA_E,'gi',weight=gSynI*0.1,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.IA_E_IA_F=Connection(self.IA_E,self.IA_F,'gi',weight=gSynI*0.1,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.IA_F_IA_E,self.IA_E_IA_F)
        
        self.R_E_IA_E=Connection(self.R_E,self.IA_E,'gi',weight=gSynI*0.1,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.R_F_IA_F=Connection(self.R_F,self.IA_F,'gi',weight=gSynI*0.1,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.IA_F_IA_E,self.IA_E_IA_F)
        
        self.MN_E_R_E=Connection(self.MN_E,self.R_E,'ge',weight=gSynE*0.25,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.MN_F_R_F=Connection(self.MN_F,self.R_F,'ge',weight=gSynE*0.25,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.MN_E_R_E,self.MN_F_R_F)
        
        self.R_F_R_E=Connection(self.R_F,self.R_E,'gi',weight=gSynI*0.1,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.R_E_R_F=Connection(self.R_E,self.R_F,'gi',weight=gSynI*0.1,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(20)*2*ms if use_delay else 0*ms);
        self.net.add(self.R_F_R_E,self.R_E_R_F)
        
        self.PF_E_MN_E=Connection(self.PF_E,self.MN_E,'ge',weight=gSynE*0.5,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(21)*2*ms if use_delay else 0*ms);
        self.PF_F_MN_F=Connection(self.PF_F,self.MN_F,'ge',weight=gSynE*0.5,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(21)*2*ms if use_delay else 0*ms);
        self.net.add(self.PF_E_MN_E,self.PF_F_MN_F)
        
        self.R_E_MN_E=Connection(self.R_E,self.MN_E,'gi',weight=gSynI*0.2,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(21)*2*ms if use_delay else 0*ms);
        self.R_F_MN_F=Connection(self.R_F,self.MN_F,'gi',weight=gSynI*0.2,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(21)*2*ms if use_delay else 0*ms);
        self.net.add(self.R_E_MN_E,self.R_F_MN_F)
        
        self.IA_F_MN_E=Connection(self.IA_F,self.MN_E,'gi',weight=gSynI*0.6,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(21)*2*ms if use_delay else 0*ms);
        self.IA_E_MN_F=Connection(self.IA_E,self.MN_F,'gi',weight=gSynI*0.6,sparseness=1,structure='dense',max_delay=10*ms if use_delay else 0 *ms,delay= lambda i,j:7*ms+randn(21)*2*ms if use_delay else 0*ms);
        self.net.add(self.IA_F_MN_E,self.IA_E_MN_F)
        
    def setup_monitors(self):
        self.traceRGE=StateMonitor(self.RG_E,'vm',record=[0])
        self.traceRGF=StateMonitor(self.RG_F,'vm',record=[0])
        self.spikesRGE=SpikeMonitor(self.RG_E)
        self.spikesRGF=SpikeMonitor(self.RG_F)
        
        self.traceINRGE=StateMonitor(self.INRG_E,'vm',record=[0])
        self.traceINRGF=StateMonitor(self.INRG_F,'vm',record=[0])
        self.spikesINRGE=SpikeMonitor(self.INRG_E)
        self.spikesINRGF=SpikeMonitor(self.INRG_F)
        
        
        self.tracePFE=StateMonitor(self.PF_E,'vm',record=[0])
        self.tracePFF=StateMonitor(self.PF_F,'vm',record=[0])
        self.spikesPFE=SpikeMonitor(self.PF_E)
        self.spikesPFF=SpikeMonitor(self.PF_F)
        
        self.traceMNE=StateMonitor(self.MN_E,'vm',record=[0])
        self.traceMNF=StateMonitor(self.MN_F,'vm',record=[0])
        self.spikesMNE=SpikeMonitor(self.MN_E)
        self.spikesMNF=SpikeMonitor(self.MN_F)
        
        self.traceIAE=StateMonitor(self.IA_E,'vm',record=[0])
        self.traceIAF=StateMonitor(self.IA_F,'vm',record=[0])
        self.spikesIAE=SpikeMonitor(self.IA_E)
        self.spikesIAF=SpikeMonitor(self.IA_F)
 
        
        self.PRRGE=PopulationRateMonitor(self.RG_E, self.binsize)
        self.PRRGF=PopulationRateMonitor(self.RG_F, self.binsize)
        
        self.PRPFE=PopulationRateMonitor(self.PF_E, self.binsize)
        self.PRPFF=PopulationRateMonitor(self.PF_F, self.binsize)
        
        self.PRMNE=PopulationRateMonitor(self.MN_E, self.binsize)
        self.PRMNF=PopulationRateMonitor(self.MN_F, self.binsize)
        
        self.net.add(self.spikesIAE,self.spikesIAF,self.traceIAE,self.traceIAF,self.PRRGE,self.PRRGF,self.PRPFE,self.PRPFF,self.PRMNE,self.PRMNF,self.traceINRGE,self.traceINRGF,self.spikesINRGE,self.spikesINRGF,self.traceRGE,self.traceRGF,self.spikesRGE,self.spikesRGF,self.spikesPFE,self.spikesPFF,self.traceMNE,self.traceMNF,self.spikesMNE,self.spikesMNF,self.tracePFE,self.tracePFF)
        
        self.Tracelist = list()
        self.Tracelist.append(self.traceRGE)
        self.Tracelist.append(self.traceRGF)
        self.Tracelist.append(self.tracePFE)
        self.Tracelist.append(self.tracePFF)
        self.Tracelist.append(self.traceIAE)
        self.Tracelist.append(self.traceIAF)
        self.Tracelist.append(self.traceMNE)
        self.Tracelist.append(self.traceMNF)
        
        self.PRlist = list()
        self.PRlist.append(self.PRRGE)
        self.PRlist.append(self.PRRGF)
        self.PRlist.append(self.PRPFE)
        self.PRlist.append(self.PRPFF)
        self.PRlist.append(self.PRMNE)
        self.PRlist.append(self.PRMNF)
        
        self.Spikelist = list()
        self.Spikelist.append(self.spikesRGE)
        self.Spikelist.append(self.spikesRGF)
        self.Spikelist.append(self.spikesPFE)
        self.Spikelist.append(self.spikesPFF)
        self.Spikelist.append(self.spikesIAE)
        self.Spikelist.append(self.spikesIAF)
        self.Spikelist.append(self.spikesMNE)
        self.Spikelist.append(self.spikesMNF)
        
        #PSP of motoneurons
        
        self.trace_MNE_ipsp=StateMonitor(self.MN_E,'ipsp',record=[0])
        self.trace_MNE_epsp=StateMonitor(self.MN_E,'epsp',record=[0])
        self.trace_MNF_ipsp=StateMonitor(self.MN_F,'ipsp',record=[0])
        self.trace_MNF_epsp=StateMonitor(self.MN_F,'epsp',record=[0])
        self.net.add(self.trace_MNE_epsp,self.trace_MNE_ipsp,self.trace_MNF_epsp,self.trace_MNF_ipsp)
        
        
        self.MN_PSPlist = list()
        self.MN_PSPlist.append(self.trace_MNE_ipsp)
        self.MN_PSPlist.append(self.trace_MNE_epsp)
        self.MN_PSPlist.append(self.trace_MNF_ipsp)
        self.MN_PSPlist.append(self.trace_MNF_epsp)
        
        #trace fake (no ion dynamics) MN
        self.trace_fakeMNE = StateMonitor(self.fakeMN_E,'vm',record=[0])
        self.trace_fakeMNF = StateMonitor(self.fakeMN_F,'vm',record=[0])
        self.net.add(self.trace_fakeMNE,self.trace_fakeMNF)
        
        self.fakeMNlist=list()
        self.fakeMNlist.append(self.trace_fakeMNE)
        self.fakeMNlist.append(self.trace_fakeMNF)
        
    def setup_network_operations(self):
        @network_operation(clock=EventClock(dt=200*ms))
        def progress(clk):
            print clk._t
            return
        
        self.net.add(progress)
        

    def __call__(self,para,value,DUR=10*second,report=0,display=False,filename='run',run=True):
        self.net.reinit(states=True)
        self.MNFactory.reinit()
        self.INFactory.reinit()
        self.reinit_monitors(self.Spikelist + self.PRlist + self.Tracelist)
        
                
        Moto.MnNeuronFactory.reinitNG(self.fakeMN_E,0)
        Moto.MnNeuronFactory.reinitNG(self.fakeMN_F,0)
        self.fakeMN_E.dyn = 0
        self.fakeMN_F.dyn = 0
 
        
        self.filename = filename
        for i in range(0,len(para)):
            self.filename = self.filename + ' - ' + para[i] +' = '+ value[i]
        self.filename = self.filename + '.pdf'
        self.DUR = DUR
        self.report = report
        self.display = display
        if run:
            self.inject_assignments(para,value)
            self.run()
            
    def inject_assignments(self,para,value):
        if len(para)==len(value):
            for i in range(0,len(para)):
                exec 'self.' + para[i] + '=' + value[i]
                print 'self.' + para[i] + '=' + value[i]
    
    def run(self):
        start = clock()
        if self.report == 0:
            self.net.run(self.DUR)
        else:
            self.net.run(self.DUR,report=self.report)
        stop = clock()
        print 'Duration: ' + str(stop-start)
        print 'Time to simulate 1 s : '+ str((stop-start)/(self.DUR/second))
        self.plot(self.display)

    def plot(self,display=False):
        self.filename=self.filename.replace('*','.')
        self.filename=self.filename.replace('/','.')
        pp = PdfPages(self.filename)
        
        fig1 = self.plot_all(1,self.plot_simple,self.Tracelist)
        fig2 = self.plot_all(2,self.plot_spike,self.Spikelist)
        fig3 = self.plot_all(3,self.plot_hist,self.PRlist)
        fig4 = self.plot_all(4,self.plot_simple,self.MN_PSPlist)
        fig5 = self.plot_all(5,self.plot_simple,self.fakeMNlist)
        
        
        
        fig1.set_size_inches(40,30)
        fig2.set_size_inches(40,30)
        fig3.set_size_inches(40,20)
        fig4.set_size_inches(40,20)
        fig5.set_size_inches(40,20)
        
        savetxt(self.filename + 'MN_E.txt',self.trace_fakeMNE.values)
        savetxt(self.filename + 'MN_F.txt',self.trace_fakeMNF.values)
        
        pp.savefig(fig1)
        pp.savefig(fig2)
        pp.savefig(fig3)
        pp.savefig(fig4)
        pp.savefig(fig5)
        pp.close()
        if display==True:
            show()
        else:
            close('all')
            
    @staticmethod     
    def showarraysmem(obj, name='',n=0):
        if n>30:
            return
        if isinstance(obj, ndarray): 
            size_MB = obj.nbytes/1024.**2 
            if size_MB>1.0: 
                print name, int(size_MB), 'MB' 
        else: 
            if hasattr(obj, '__dict__'): 
                for x, y in obj.__dict__.iteritems(): 
                    if len(name): 
                        namex = name+'.'+x 
                    else: 
                        namex = x 
                    rybak.showarraysmem(y, namex,n+1) 


